import pandas as pd
from sqlalchemy import create_engine
# 随机森林回归模型
from sklearn.ensemble import RandomForestRegressor
# 拆分数据集（训练集和测试集）
from sklearn.model_selection import train_test_split
# 特征编码（独热编码）
from sklearn.preprocessing import OneHotEncoder

# 读取数据
engine = create_engine('mysql+pymysql://root:123456@localhost:3306/bossinfo')
sql = 'select * from jobinfo'
data = pd.read_sql(sql, engine)

# 数据预处理
data['salary'] = data['salary'].apply(lambda x: [int(i) for i in x.strip('[]').split(',')])
data.drop_duplicates(subset='id', inplace=True)
data.fillna(value=0, inplace=True)

# 选择特征列
features = ['address', 'educational', 'type', 'workExperience']

# 提取特征和目标变量
X = data[features]
y = data['salary']

# 使用独热编码对分类特征进行编码
onehot_encoder = OneHotEncoder()
X_encoded = onehot_encoder.fit_transform(X)

# 将薪资区间转换为单个数字
y_mean = y.apply(lambda x: (x[0] + x[1]) / 2)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_encoded, y_mean, test_size=0.2, random_state=42)

# 初始化随机森林回归模型
rf_model = RandomForestRegressor(n_estimators=20)

# 训练模型
rf_model.fit(X_train, y_train)

# 获取新数据并进行预测
new_data = pd.DataFrame({'address': ['北京'], 'educational': ['本科'], 'type': ['人工智能'], 'workExperience': ['10年以上']})
new_data_encoded = onehot_encoder.transform(new_data)
prediction = rf_model.predict(new_data_encoded)
print('预测薪资：', prediction)
# 关闭连接
engine.dispose()
